# tcp_router.py
import json
from typing import Optional

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Request, Body, Query, status

from tcp_server.tcp_clients import ClientManager
from tcp_server.websocket_manager import websocket_manager
from utils.log_util import logger
from module_task.search_scheduler import search_scheduler, TaskConfig, scheduler_manager

tcpRouter = APIRouter(tags=['TCP通信'])


@tcpRouter.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket连接端点 - 用于前端实时通信"""
    # ✅ 从查询参数或headers获取token，提取user_id
    user_id = None
    token = None
    try:
        # 尝试从查询参数获取token
        token = websocket.query_params.get("token")
        logger.debug(f"🔍 从query_params获取token: {token[:20] + '...' if token and len(token) > 20 else token}")
        
        if not token:
            # 尝试从headers获取
            auth_header = websocket.headers.get("Authorization") or websocket.headers.get("authorization")
            logger.debug(f"🔍 从headers获取Authorization: {auth_header[:50] + '...' if auth_header and len(auth_header) > 50 else auth_header}")
            if auth_header:
                if auth_header.startswith("Bearer "):
                    token = auth_header[7:]
                else:
                    token = auth_header
        
        # ✅ 统一验证逻辑（包含无Token情况）
        from tcp_server.ws_auth import verify_websocket_token
        is_valid, result = verify_websocket_token(token)
        logger.debug(f"🔍 Token验证结果: is_valid={is_valid}, result={result}")
        
        if not is_valid:
            logger.warning(f"⚠️ WebSocket认证失败: {result}，拒绝连接")
            # 拒绝连接
            await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
            return

        # ✅ 认证通过，提取user_id
        if result == "guest":
            # 访客模式或未开启认证
            logger.warning(f"⚠️ WebSocket以访客模式连接: {result}")
            user_id = None
        else:
            try:
                user_id = int(result)
                logger.debug(f"✅ WebSocket连接已认证: user_id={user_id}")
            except (ValueError, TypeError) as e:
                logger.warning(f"⚠️ 无法将result转换为user_id: result={result}, error={e}")
                user_id = None
                
    except Exception as e:
        logger.warning(f"⚠️ 获取WebSocket user_id失败: {e}", exc_info=True)
    
    # websocket.accept() 由 websocket_manager.connect() 内部调用
    # ✅ 记录连接前的状态
    logger.info(f"🔌 WebSocket连接请求: token存在={bool(token)}, user_id={user_id}, query_params={dict(websocket.query_params)}, headers.Authorization={websocket.headers.get('Authorization', 'N/A')[:50] if websocket.headers.get('Authorization') else 'N/A'}")
    
    websocket_id = await websocket_manager.connect(websocket, user_id=user_id)
    
    # ✅ 记录连接后的状态
    if user_id is None:
        logger.warning(
            f"⚠️ WebSocket连接：user_id为None，可能导致消息无法正确推送。"
            f"websocket_id={websocket_id}, "
            f"当前websocket_user_map={dict(websocket_manager.websocket_user_map)}, "
            f"当前连接数={len(websocket_manager.active_connections)}"
        )
    else:
        logger.info(
            f"✅ WebSocket客户端连接: websocket_id={websocket_id}, user_id={user_id}, "
            f"websocket_user_map已更新={dict(websocket_manager.websocket_user_map)}"
        )
    
    try:
        # 发送当前在线客户端列表（从数据库获取完整信息，包含中文状态）
        from config.get_db import get_db
        from module_admin.service.agent_client_service import AgentClientService
        from utils.status_translation import translate_client_info
        
        online_clients = []
        async for db in get_db():
            clients = await AgentClientService.get_online_clients(db)
            online_clients = [
                translate_client_info({
                    "id": client.id,
                    "client_id": client.client_id,
                    "client_name": client.client_name,
                    "ip_address": client.ip_address,
                    "port": client.port,
                    "hostname": client.hostname or "",  # 🔥 添加主机名
                    "client_type": client.client_type,
                    "online_status": client.online_status,
                    "work_status": client.work_status,
                    "browser_status": client.browser_status,
                    "login_status": client.login_status,
                    "account": client.account or "",
                    "account_nickname": client.account_nickname or "",  # 🔥 添加昵称
                    "default_address_id": getattr(client, 'default_address_id', None) or "",
                    "remark": client.remark or "",
                    "connected_at": client.connected_at.isoformat() if client.connected_at else "",
                    "last_heartbeat": client.last_heartbeat.isoformat() if client.last_heartbeat else "",
                })
                for client in clients
            ]
            break  # 只需要第一次查询结果
        
        await websocket.send_json({
            "type": "init",
            "clients": online_clients
        })
        
        # 接收消息循环
        while True:
            data = await websocket.receive_text()
            message = json.loads(data)
            
            # 处理前端发来的命令
            await handle_websocket_message(websocket, message)
            
    except WebSocketDisconnect:
        logger.info(f"WebSocket客户端断开: {websocket_id}")
    except Exception as e:
        logger.exception(f"WebSocket错误: {e}")
    finally:
        websocket_manager.disconnect(websocket_id)


async def handle_websocket_message(websocket: WebSocket, message: dict):
    """处理WebSocket消息"""
    msg_type = message.get("type")
    event = message.get("event")
    direction = message.get("direction")

    # ✅ 处理前端通过sendCommand发送的命令（type="command"格式）
    if msg_type == "command":
        command = message.get("command")
        client_id = message.get("client_id")
        data = message.get("data", {})
        
        if not command:
            await websocket.send_json({
                "type": "error",
                "message": "缺少command参数"
            })
            return
        
        if not client_id:
            await websocket.send_json({
                "type": "error",
                "message": "缺少client_id参数"
            })
            return
        
        # 构造符合Agent格式的命令消息
        import uuid
        cmd_message = {
            "event": command,
            "direction": "req",
            "success": True,
            "code": 0,
            "msg": "",
            "data": data,
            "req_id": str(uuid.uuid4())
        }
        
        logger.info(f"📤 转发命令到Agent: client_id={client_id}, command={command}, data={data}")
        
        # 发送到TCP客户端
        success = await ClientManager.send_to_client(client_id, cmd_message)
        
        # 返回结果
        await websocket.send_json({
            "type": "command_result",
            "success": success,
            "client_id": client_id,
            "command": command
        })
        return
    
    # 处理刷新客户端列表请求
    if event == "refresh_clients" and direction == "req":
        from config.get_db import get_db
        from module_admin.service.agent_client_service import AgentClientService
        from utils.status_translation import translate_client_info
        
        online_clients = []
        async for db in get_db():
            clients = await AgentClientService.get_online_clients(db)
            online_clients = [
                translate_client_info({
                    "id": client.id,
                    "client_id": client.client_id,
                    "client_name": client.client_name,
                    "ip_address": client.ip_address,
                    "port": client.port,
                    "hostname": client.hostname or "",  # 🔥 添加主机名
                    "client_type": client.client_type,
                    "online_status": client.online_status,
                    "work_status": client.work_status,
                    "browser_status": client.browser_status,
                    "login_status": client.login_status,
                    "account": client.account or "",
                    "account_nickname": client.account_nickname or "",  # 🔥 添加昵称
                    "default_address_id": getattr(client, 'default_address_id', None) or "",
                    "remark": client.remark or "",
                    "connected_at": client.connected_at.isoformat() if client.connected_at else "",
                    "last_heartbeat": client.last_heartbeat.isoformat() if client.last_heartbeat else "",
                })
                for client in clients
            ]
            break
        
        await websocket.send_json({
            "type": "init",
            "clients": online_clients
        })
        return
    
    # 处理前端发送的命令（event + direction 格式）
    if event and direction == "req":
        import uuid
        
        # 从data中提取client_id
        data = message.get("data", {})
        client_id = data.get("client_id")
        
        if not client_id:
            logger.warning(f"收到命令但缺少client_id: event={event}")
            await websocket.send_json({
                "type": "error",
                "message": "缺少client_id参数"
            })
            return
        
        # 构造符合Agent格式的命令消息
        cmd_message = {
            "event": event,
            "direction": "req",
            "success": True,
            "code": 0,
            "msg": "",
            "data": {k: v for k, v in data.items() if k != "client_id"},
            "req_id": str(uuid.uuid4())
        }
        
        logger.info(f"📤 转发命令到Agent: client_id={client_id}, event={event}, data={data}")
        
        # 发送到TCP客户端
        success = await ClientManager.send_to_client(client_id, cmd_message)
        
        if not success:
            # 打印更详细的错误信息
            all_clients = [c.client_id for c in ClientManager.get_all_clients()]
            logger.error(f"❌ 发送命令失败: client_id={client_id}, event={event}")
            logger.error(f"当前在线客户端: {all_clients}")
        else:
            logger.success(f"✅ 命令已发送到Agent: client_id={client_id}, event={event}")
        
        # 返回结果
        await websocket.send_json({
            "type": "command_result",
            "success": success,
            "client_id": client_id,
            "event": event
        })
        return
    
    # 处理心跳
    if msg_type == "ping":
        # 心跳响应
        await websocket.send_json({"type": "pong"})


@tcpRouter.get("/clients/list")
async def get_tcp_clients():
    """获取TCP客户端列表（HTTP接口）"""
    clients = [
        {
            "client_id": client.client_id,
            "ip_address": client.ip_address,
            "port": client.port,
            "connected_at": client.connected_at.isoformat(),
            "last_heartbeat": client.last_heartbeat.isoformat(),
            "status": client.status,
            "current_command": client.current_command
        }
        for client in ClientManager.get_all_clients()
    ]
    
    return {
        "code": 200,
        "msg": "success",
        "data": clients
    }


@tcpRouter.post("/clients/{client_id}/command")
async def send_command_to_client(client_id: str, command_data: dict):
    """向指定客户端发送命令（HTTP接口）"""
    import uuid
    
    command = command_data.get("command")
    data = command_data.get("data", {})
    
    if not command:
        return {
            "code": 400,
            "msg": "缺少command参数"
        }
    
    # 构造符合Agent格式的命令消息
    cmd_message = {
        "event": command,
        "direction": "req",
        "success": True,
        "code": 0,
        "msg": "",
        "data": data,
        "req_id": str(uuid.uuid4())
    }
    
    # 发送到TCP客户端
    success = await ClientManager.send_to_client(client_id, cmd_message)
    
    if success:
        return {
            "code": 200,
            "msg": "命令已发送",
            "data": {"client_id": client_id, "command": command}
        }
    else:
        return {
            "code": 500,
            "msg": "发送失败，客户端可能已离线"
        }


# ============== 搜索任务调度器控制 ==============
@tcpRouter.post("/scheduler/start")
async def scheduler_start(config: dict = Body(...), request: Request = None):
    """启动搜索调度器（支持多用户独立调度）"""
    try:
        # ✅ 从请求中获取user_id（如果前端传递了）
        user_id = config.get("user_id")
        if not user_id and request:
            user_id = getattr(request.state, "user_id", None)
        task_id = config.get("task_id")
        
        cfg = TaskConfig(
            keyword=config.get("keyword", ""),
            total_time=int(config.get("total_time", 60) or 60),
            enable_auto_order=bool(config.get("enable_auto_order", False)),
            skip_first_search=bool(config.get("skip_first_search", False)),  # ✅ 首次搜索不自动下单
            order_price=float(config.get("order_price", 0) or 0),
            min_price=float(config.get("min_price", 0) or 0),
            max_price=float(config.get("max_price", 0) or 0),
            display_count=int(config.get("display_count", 30) or 30),
            user_id=user_id,  # ✅ 传递user_id
            task_id=task_id   # ✅ 传递task_id
        )
        if not cfg.keyword:
            return {"code": 400, "msg": "keyword不能为空"}
        
        # ✅ 如果提供了user_id，使用调度器管理器（多用户支持）
        if user_id:
            scheduler = await scheduler_manager.get_scheduler(user_id)
            await scheduler.start(cfg)
            logger.info(f"✅ 用户 {user_id} 的调度器已启动（任务ID={task_id}，关键词={cfg.keyword}）")
        else:
            # 兼容旧代码：使用全局单例（不推荐）
            await search_scheduler.start(cfg)
            logger.warning("⚠️ 使用全局单例调度器（未提供user_id），建议传递user_id以支持多用户隔离")
        
        return {"code": 200, "msg": "scheduler started"}
    except Exception as e:
        logger.exception(e)
        return {"code": 500, "msg": str(e)}


@tcpRouter.post("/scheduler/stop")
async def scheduler_stop(config: dict = Body(default=None), request: Request = None):
    """停止搜索调度器（支持多用户）"""
    try:
        # ✅ 如果提供了user_id，停止指定用户的调度器
        user_id = None
        if config:
            user_id = config.get("user_id")
        if not user_id and request:
            user_id = getattr(request.state, "user_id", None)
        
        if user_id:
            scheduler = await scheduler_manager.get_scheduler(user_id)
            await scheduler.stop()
            logger.info(f"✅ 用户 {user_id} 的调度器已停止")
        else:
            # 兼容旧代码：停止全局单例
            await search_scheduler.stop()
            logger.warning("⚠️ 停止全局单例调度器（未提供user_id）")
        
        return {"code": 200, "msg": "scheduler stopped"}
    except Exception as e:
        logger.exception(e)
        return {"code": 500, "msg": str(e)}


@tcpRouter.get("/scheduler/status")
async def scheduler_status(
    request: Request, 
    user_id: Optional[int] = Query(None, description="用户ID（可选，用于多用户隔离）")
):
    """获取搜索调度器状态（包含统计信息）
    
    Args:
        user_id: 用户ID（可选，用于多用户隔离）。如果不提供，从JWT token中获取当前用户ID
    """
    # ✅ 从请求中获取user_id（优先使用查询参数，否则从JWT token获取）
    if user_id is None:
        user_id = getattr(request.state, "user_id", None)
        if user_id:
            user_id = int(user_id)
    
    # ✅ 使用scheduler_manager获取对应用户的调度器
    if user_id:
        scheduler = await scheduler_manager.get_scheduler(user_id)
    else:
        # 兼容旧代码：如果没有user_id，使用默认调度器（不推荐）
        logger.warning("scheduler_status: 未提供user_id，使用默认调度器（可能不准确）")
        scheduler = search_scheduler
    
    cfg = scheduler.get_config()
    is_running = scheduler.is_running()
    
    # 获取可用客户端数量
    search_client_count = 0
    order_client_count = 0
    current_client_id = None
    next_search_time = None
    
    if is_running and cfg:
        try:
            from config.get_db import get_db
            from module_admin.service.agent_client_service import AgentClientService
            
            async for db in get_db():
                # ✅ 获取可用搜索客户端（按user_id过滤）
                search_clients = await AgentClientService.get_available_search_clients(
                    db, 
                    ignore_online_status=True,
                    user_id=user_id  # ✅ 按user_id过滤
                )
                # ✅ 过滤条件：在线 && 没有被封禁 && 不是异常状态
                online_search = [
                    c for c in search_clients 
                    if c.online_status == 'online'
                    and (c.work_status or 'idle') != 'banned'  # ✅ 排除被封禁的账户
                    and (c.work_status or 'idle') != 'error'   # ✅ 排除异常状态
                ]
                search_client_count = len(online_search)
                
                # ✅ 获取可用下单客户端（按user_id过滤）
                order_clients = await AgentClientService.get_available_order_clients(
                    db,
                    user_id=user_id  # ✅ 按user_id过滤
                )
                online_order = [c for c in order_clients if c.online_status == 'online']
                order_client_count = len(online_order)
                
                # 获取当前正在搜索的客户端（从调度器）
                if is_running:
                    current_client_id = await scheduler.get_current_client_id()
                
                # 如果有运行的客户端，计算下次搜索时间（格式化显示）
                if search_client_count > 0 and cfg.total_time:
                    from datetime import datetime, timedelta
                    import time
                    # 计算下一轮搜索的预估时间（基于总时间和客户端数量）
                    slice_time = max(cfg.total_time // search_client_count, 1)
                    next_timestamp = datetime.now() + timedelta(seconds=slice_time)
                    next_search_time = next_timestamp.strftime('%H:%M:%S')
                
                break
        except Exception as e:
            logger.warning(f"获取客户端统计失败: {e}")
    
    # ✅ 获取轮次计数
    round_count = 0
    if scheduler:
        round_count = scheduler.get_round_count()
    
    return {
        "code": 200,
        "msg": "success",
        "data": {
            "running": is_running,
            "config": cfg.__dict__ if cfg else None,
            "search_client_count": search_client_count,
            "order_client_count": order_client_count,
            "current_client_id": current_client_id,
            "next_search_time": next_search_time,
            "round_count": round_count  # ✅ 轮次计数
        }
    }
